import gym
import torch
import torch.nn as nn
import torch.optim as optim
import torch.nn.functional as F
from torch.distributions import Categorical, Normal
import numpy as np
import matplotlib.pyplot as plt
from collections import deque
import random


# REINFORCE 算法
class ReinforceAgent:
    def __init__(self, state_dim, action_dim, gamma=0.99, lr=0.001):
        self.policy = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, action_dim),
            nn.Softmax(dim=-1)
        )
        self.optimizer = optim.Adam(self.policy.parameters(), lr=lr)
        self.gamma = gamma
        self.log_probs = []
        self.rewards = []

    def select_action(self, state):
        state = torch.FloatTensor(state)
        probs = self.policy(state)
        m = Categorical(probs)
        action = m.sample()
        self.log_probs.append(m.log_prob(action))
        return action.item()

    def update(self):
        R = 0
        policy_loss = []
        returns = deque()
        for r in reversed(self.rewards):
            R = r + self.gamma * R
            returns.appendleft(R)
        returns = torch.tensor(returns)
        returns = (returns - returns.mean()) / (returns.std() + 1e-8)
        for log_prob, R in zip(self.log_probs, returns):
            policy_loss.append(-log_prob * R)
        self.optimizer.zero_grad()
        torch.cat(policy_loss).sum().backward()
        self.optimizer.step()
        self.log_probs = []
        self.rewards = []


# Actor - Critic 算法
class ActorCriticAgent:
    def __init__(self, state_dim, action_dim, gamma=0.99, lr=0.001):
        self.actor = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, action_dim),
            nn.Softmax(dim=-1)
        )
        self.critic = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=lr)
        self.gamma = gamma
        self.log_probs = []
        self.rewards = []
        self.states = []

    def select_action(self, state):
        state = torch.FloatTensor(state)
        probs = self.actor(state)
        m = Categorical(probs)
        action = m.sample()
        self.log_probs.append(m.log_prob(action))
        self.states.append(state)
        return action.item()

    def update(self):
        R = 0
        returns = []
        for r in reversed(self.rewards):
            R = r + self.gamma * R
            returns.insert(0, R)
        returns = torch.tensor(returns)
        returns = (returns - returns.mean()) / (returns.std() + 1e-8)
        states = torch.stack(self.states)
        values = self.critic(states).detach()
        advantages = returns - values
        actor_loss = (-torch.stack(self.log_probs) * advantages).sum()
        critic_loss = F.mse_loss(self.critic(states), returns)
        self.actor_optimizer.zero_grad()
        self.critic_optimizer.zero_grad()
        actor_loss.backward()
        critic_loss.backward()
        self.actor_optimizer.step()
        self.critic_optimizer.step()
        self.log_probs = []
        self.rewards = []
        self.states = []


# DDPG 算法
class DDPG:
    def __init__(self, state_dim, action_dim, gamma=0.99, lr=0.001):
        self.actor = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, action_dim),
            nn.Tanh()
        )
        self.critic = nn.Sequential(
            nn.Linear(state_dim + action_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )
        self.target_actor = nn.Sequential(*list(self.actor.children()))
        self.target_critic = nn.Sequential(*list(self.critic.children()))
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=lr)
        self.gamma = gamma
        self.memory = []
        self.batch_size = 64
        self.tau = 0.001

    def select_action(self, state):
        state = torch.FloatTensor(state).unsqueeze(0)
        return self.actor(state).detach().numpy()[0]

    def store_transition(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))
        if len(self.memory) > 10000:
            self.memory.pop(0)

    def update(self):
        if len(self.memory) < self.batch_size:
            return
        batch = random.sample(self.memory, self.batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        states = torch.FloatTensor(states)
        actions = torch.FloatTensor(actions)
        rewards = torch.FloatTensor(rewards).unsqueeze(1)
        next_states = torch.FloatTensor(next_states)
        dones = torch.FloatTensor(dones).unsqueeze(1)

        target_actions = self.target_actor(next_states)
        target_Q = self.target_critic(torch.cat([next_states, target_actions], 1))
        target_Q = rewards + (1 - dones) * self.gamma * target_Q
        current_Q = self.critic(torch.cat([states, actions], 1))
        critic_loss = F.mse_loss(current_Q, target_Q.detach())

        actor_loss = -self.critic(torch.cat([states, self.actor(states)], 1)).mean()

        self.actor_optimizer.zero_grad()
        self.critic_optimizer.zero_grad()
        actor_loss.backward()
        critic_loss.backward()
        self.actor_optimizer.step()
        self.critic_optimizer.step()

        for target_param, param in zip(self.target_actor.parameters(), self.actor.parameters()):
            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
        for target_param, param in zip(self.target_critic.parameters(), self.critic.parameters()):
            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)


# PPO 算法
class PPO:
    def __init__(self, state_dim, action_dim, gamma=0.99, lr=0.001, clip_ratio=0.2):
        self.actor = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, action_dim),
            nn.Softmax(dim=-1)
        )
        self.critic = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr)
        self.critic_optimizer = optim.Adam(self.critic.parameters(), lr=lr)
        self.gamma = gamma
        self.clip_ratio = clip_ratio
        self.log_probs = []
        self.rewards = []
        self.states = []
        self.actions = []

    def select_action(self, state):
        state = torch.FloatTensor(state)
        probs = self.actor(state)
        m = Categorical(probs)
        action = m.sample()
        self.log_probs.append(m.log_prob(action))
        self.states.append(state)
        self.actions.append(action)
        return action.item()

    def update(self):
        R = 0
        returns = []
        for r in reversed(self.rewards):
            R = r + self.gamma * R
            returns.insert(0, R)
        returns = torch.tensor(returns)
        returns = (returns - returns.mean()) / (returns.std() + 1e-8)
        states = torch.stack(self.states)
        values = self.critic(states).detach()
        advantages = returns - values
        probs = self.actor(states)
        m = Categorical(probs)
        new_log_probs = m.log_prob(torch.tensor(self.actions))
        ratio = torch.exp(new_log_probs - torch.stack(self.log_probs))
        surr1 = ratio * advantages
        surr2 = torch.clamp(ratio, 1 - self.clip_ratio, 1 + self.clip_ratio) * advantages
        actor_loss = (-torch.min(surr1, surr2)).mean()
        critic_loss = F.mse_loss(self.critic(states), returns)
        self.actor_optimizer.zero_grad()
        self.critic_optimizer.zero_grad()
        actor_loss.backward()
        critic_loss.backward()
        self.actor_optimizer.step()
        self.critic_optimizer.step()
        self.log_probs = []
        self.rewards = []
        self.states = []
        self.actions = []


# SAC 算法
class SAC:
    def __init__(self, state_dim, action_dim, gamma=0.99, lr=0.001, alpha=0.2):
        self.actor = nn.Sequential(
            nn.Linear(state_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 2 * action_dim)
        )
        self.critic1 = nn.Sequential(
            nn.Linear(state_dim + action_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )
        self.critic2 = nn.Sequential(
            nn.Linear(state_dim + action_dim, 64),
            nn.ReLU(),
            nn.Linear(64, 1)
        )
        self.target_critic1 = nn.Sequential(*list(self.critic1.children()))
        self.target_critic2 = nn.Sequential(*list(self.critic2.children()))
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=lr)
        self.critic1_optimizer = optim.Adam(self.critic1.parameters(), lr=lr)
        self.critic2_optimizer = optim.Adam(self.critic2.parameters(), lr=lr)
        self.gamma = gamma
        self.alpha = alpha
        self.memory = []
        self.batch_size = 64
        self.tau = 0.001

    def select_action(self, state):
        state = torch.FloatTensor(state).unsqueeze(0)
        mu, log_std = self.actor(state).chunk(2, dim=-1)
        std = log_std.exp()
        m = Normal(mu, std)
        action = m.rsample()
        return torch.tanh(action).detach().numpy()[0]

    def store_transition(self, state, action, reward, next_state, done):
        self.memory.append((state, action, reward, next_state, done))
        if len(self.memory) > 10000:
            self.memory.pop(0)

    def update(self):
        if len(self.memory) < self.batch_size:
            return
        batch = random.sample(self.memory, self.batch_size)
        states, actions, rewards, next_states, dones = zip(*batch)
        states = torch.FloatTensor(states)
        actions = torch.FloatTensor(actions)
        rewards = torch.FloatTensor(rewards).unsqueeze(1)
        next_states = torch.FloatTensor(next_states)
        dones = torch.FloatTensor(dones).unsqueeze(1)

        mu, log_std = self.actor(states).chunk(2, dim=-1)
        std = log_std.exp()
        m = Normal(mu, std)
        new_actions = m.rsample()
        log_prob = m.log_prob(new_actions).sum(-1, keepdim=True)
        new_actions = torch.tanh(new_actions)

        q1 = self.critic1(torch.cat([states, new_actions], 1))
        q2 = self.critic2(torch.cat([states, new_actions], 1))
        q = torch.min(q1, q2)
        actor_loss = (self.alpha * log_prob - q).mean()

        target_mu, target_log_std = self.actor(next_states).chunk(2, dim=-1)
        target_std = target_log_std.exp()
        target_m = Normal(target_mu, target_std)
        target_new_actions = target_m.rsample()
        target_log_prob = target_m.log_prob(target_new_actions).sum(-1, keepdim=True)
        target_new_actions = torch.tanh(target_new_actions)

        target_q1 = self.target_critic1(torch.cat([next_states, target_new_actions], 1))
        target_q2 = self.target_critic2(torch.cat([next_states, target_new_actions], 1))
        target_q = torch.min(target_q1, target_q2)
        target_q = rewards + (1 - dones) * self.gamma * (target_q - self.alpha * target_log_prob)

        current_q1 = self.critic1(torch.cat([states, actions], 1))
        current_q2 = self.critic2(torch.cat([states, actions], 1))
        critic1_loss = F.mse_loss(current_q1, target_q.detach())
        critic2_loss = F.mse_loss(current_q2, target_q.detach())

        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        self.actor_optimizer.step()

        self.critic1_optimizer.zero_grad()
        critic1_loss.backward()
        self.critic1_optimizer.step()

        self.critic2_optimizer.zero_grad()
        critic2_loss.backward()
        self.critic2_optimizer.step()

        for target_param, param in zip(self.target_critic1.parameters(), self.critic1.parameters()):
            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)
        for target_param, param in zip(self.target_critic2.parameters(), self.critic2.parameters()):
            target_param.data.copy_(self.tau * param.data + (1 - self.tau) * target_param.data)


# 训练函数
def train(agent, env, episodes=1000):
    all_rewards = []
    for episode in range(episodes):
        state = env.reset()
        episode_reward = 0
        done = False
        while not done:
            action = agent.select_action(state)
            next_state, reward, done, _ = env.step(action)
            if isinstance(agent, (DDPG, SAC)):
                agent.store_transition(state, action, reward, next_state, done)
            else:
                agent.rewards.append(reward)
            state = next_state
            episode_reward += reward
            if isinstance(agent, (DDPG, SAC)):
                agent.update()
        if not isinstance(agent, (DDPG, SAC)):
            agent.update()
        all_rewards.append(episode_reward)
        if episode % 100 == 0:
            print(f"Episode {episode}: Reward = {episode_reward}")
    return all_rewards


# 主程序
if __name__ == "__main__":
    env = gym.make('CartPole-v0')
    state_dim = env.observation_space.shape[0]
    action_dim = env.action_space.n

    # REINFORCE
    reinforce_agent = ReinforceAgent(state_dim, action_dim)
    reinforce_rewards = train(reinforce_agent, env)

    # Actor - Critic
    ac_agent = ActorCriticAgent(state_dim, action_dim)
    ac_rewards = train(ac_agent, env)

    # DDPG
    ddpg_agent = DDPG(state_dim, action_dim)
    ddpg_rewards = train(ddpg_agent, env)

    # PPO
    ppo_agent = PPO(state_dim, action_dim)
    ppo_rewards = train(ppo_agent, env)

    # SAC
    sac_agent = SAC(state_dim, action_dim)
    sac_rewards = train(sac_agent, env)

    # 可视化结果
    plt.figure(figsize=(12, 8))
    plt.plot(reinforce_rewards, label='REINFORCE')
    plt.plot(ac_rewards, label='Actor - Critic')
    plt.plot(ddpg_rewards, label='DDPG')
    plt.plot(ppo_rewards, label='PPO')
    plt.plot(sac_rewards, label='SAC')
    plt.xlabel('Episodes')
    plt.ylabel('Reward')
    plt.title('Training Rewards of Different Algorithms')
    plt.legend()
    plt.show()

    